"""
TODO:
    不识别InstanceNorm
"""
import openvino as ov
import time
import cv2
import numpy as np

core = ov.Core()
cpu_model = core.read_model("v1_openvino.xml")
model = core.compile_model(model=cpu_model)

def torch_style_softmax_numpy_impl(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

def openvino_infer_data(x):
    return model([mat])["output:0"]

def get_nnunet_predict(x: np.ndarray, dims: list):
    return np.flip(
            torch_style_softmax_numpy_impl(openvino_infer_data(
                np.flip(x, dims)
            ), 1),
            dims
    )

sample_image = "../resource/raw_resize/1.jpg"
start_time = time.perf_counter()
mat = cv2.imread(sample_image)
mat = mat.transpose((2, 0, 1))
mat = mat.astype(np.float32)
for c in range(mat.shape[0]):
    mat[c] = (mat[c] - mat[c].mean()) / (mat[c].std() + 1e-8)
mat = np.expand_dims(mat, 0)

pred = torch_style_softmax_numpy_impl(openvino_infer_data(mat), 1)
pred += get_nnunet_predict(mat, [3])
pred += get_nnunet_predict(mat, [2])
pred += get_nnunet_predict(mat, [3, 2])
pred /= 4.0

# save
pred = np.argmax(pred, 1)
pred = pred[0]
cv2.imwrite("out_pred_openvino.png", pred * 255)

end_time = time.perf_counter()
print("time cost: ", end_time - start_time, "s")